from sentence_transformers import SentenceTransformer
import os
from config import LOCAL_EMBEDDING_CONFIG
import torch

class LocalEmbedding:
    def __init__(self, model_path=None, device=None, normalize_embeddings=None):
        """
        初始化本地embedding模型
        Args:
            model_path: 本地模型路径，如果为None则使用配置文件中的路径
            device: 运行设备，如果为None则使用配置文件中的设备
            normalize_embeddings: 是否对embedding进行归一化，如果为None则使用配置文件中的设置
        """
        self.model_path = model_path or LOCAL_EMBEDDING_CONFIG['model_path']
        self.device = device or LOCAL_EMBEDDING_CONFIG['device']
        self.normalize_embeddings = normalize_embeddings if normalize_embeddings is not None else LOCAL_EMBEDDING_CONFIG['normalize_embeddings']
        
        if not os.path.exists(self.model_path):
            raise ValueError(f"模型路径 {self.model_path} 不存在，请确保已下载模型文件")
            
        # 加载模型到指定设备
        self.model = SentenceTransformer(self.model_path)
        self.model.to(self.device)
    
    def get_embeddings(self, texts):
        """
        获取文本的embedding向量
        Args:
            texts: 输入文本或文本列表
        Returns:
            List[List[float]]: embedding向量列表
        """
        if isinstance(texts, str):
            texts = [texts]
        
        # 使用模型生成embedding
        with torch.no_grad():
            embeddings = self.model.encode(
                texts,
                normalize_embeddings=self.normalize_embeddings,
                convert_to_tensor=True,
                device=self.device
            )
            
        # 转换为列表格式
        return embeddings.cpu().numpy().tolist()

def create_embedding_instance(model_path=None, device=None, normalize_embeddings=None):
    """
    创建embedding实例的工厂函数
    Args:
        model_path: 可选的本地模型路径
        device: 可选的运行设备
        normalize_embeddings: 可选的是否归一化设置
    Returns:
        LocalEmbedding: 本地embedding实例
    """
    return LocalEmbedding(
        model_path=model_path,
        device=device,
        normalize_embeddings=normalize_embeddings
    ) 